package org.hibernate.tool.hbm2ddl;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.Reader;
import java.io.Writer;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.SQLWarning;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.hibernate.HibernateException;
import org.hibernate.JDBCException;
import org.hibernate.cfg.Configuration;
import org.hibernate.cfg.Environment;
import org.hibernate.cfg.MyConfiguration;
import org.hibernate.cfg.NamingStrategy;
import org.hibernate.cfg.Settings;
import org.hibernate.dialect.Dialect;
import org.hibernate.mapping.PersistentClass;
import org.hibernate.pretty.DDLFormatter;
import org.hibernate.util.ConfigHelper;
import org.hibernate.util.JDBCExceptionReporter;
import org.hibernate.util.PropertiesHelper;
import org.hibernate.util.ReflectHelper;

public class MySchemaExport {

	private static final Log log = LogFactory.getLog(SchemaExport.class);

	private ConnectionHelper connectionHelper;
	private String[] dropSQL;
	private String[] createSQL;
	private String outputFile = null;
	private String importFile = "/import.sql";
	private Dialect dialect;
	private String delimiter;
	private final List exceptions = new ArrayList();
	private boolean haltOnError = false;
	private boolean format = true;

	/**
	 * Create a schema exporter for the given Configuration and given settings
	 */
	public MySchemaExport(MyConfiguration cfg, Settings settings,
			PersistentClass[] exports) throws HibernateException {
		dialect = settings.getDialect();
		connectionHelper = new SuppliedConnectionProviderConnectionHelper(
				settings.getConnectionProvider());
		dropSQL = cfg.generateDropSchemaScript(dialect, exports);
		createSQL = cfg.generateSchemaCreationScript(dialect, exports);
		format = settings.isFormatSqlEnabled();
	}

	/**
	 * Run the schema creation script.
	 * 
	 * @param script
	 *            print the DDL to the console
	 * @param export
	 *            export the script to the database
	 */
	public void create(boolean script, boolean export) {
		execute(script, export, false, false);
	}

	/**
	 * Run the drop schema script.
	 * 
	 * @param script
	 *            print the DDL to the console
	 * @param export
	 *            export the script to the database
	 */
	public void drop(boolean script, boolean export) {
		execute(script, export, true, false);
	}

	private String format(String sql) {
		return format ? new DDLFormatter(sql).format() : sql;
	}

	public void execute(boolean script, boolean export, boolean justDrop,
			boolean justCreate) {

		log.info("Running hbm2ddl schema export");

		Connection connection = null;
		Writer outputFileWriter = null;
		Reader importFileReader = null;
		Statement statement = null;

		exceptions.clear();

		try {

			try {
				InputStream stream = ConfigHelper
						.getResourceAsStream(importFile);
				importFileReader = new InputStreamReader(stream);
			} catch (HibernateException e) {
				log.debug("import file not found: " + importFile);
			}

			if (outputFile != null) {
				log.info("writing generated schema to file: " + outputFile);
				outputFileWriter = new FileWriter(outputFile);
			}

			if (export) {
				log.info("exporting generated schema to database");
				connectionHelper.prepare(true);
				connection = connectionHelper.getConnection();
				statement = connection.createStatement();
			}

			if (!justCreate) {
				drop(script, export, outputFileWriter, statement);
			}

			if (!justDrop) {
				create(script, export, outputFileWriter, statement);
				if (export && importFileReader != null) {
					importScript(importFileReader, statement);
				}
			}

			log.info("schema export complete");

		}

		catch (Exception e) {
			exceptions.add(e);
			log.error("schema export unsuccessful", e);
		}

		finally {

			try {
				if (statement != null) {
					statement.close();
				}
				if (connection != null) {
					connectionHelper.release();
				}
			} catch (Exception e) {
				exceptions.add(e);
				log.error("Could not close connection", e);
			}

			try {
				if (outputFileWriter != null) {
					outputFileWriter.close();
				}
				if (importFileReader != null) {
					importFileReader.close();
				}
			} catch (IOException ioe) {
				exceptions.add(ioe);
				log.error("Error closing output file: " + outputFile, ioe);
			}

		}
	}

	private void importScript(Reader importFileReader, Statement statement)
			throws IOException {
		log.info("Executing import script: " + importFile);
		BufferedReader reader = new BufferedReader(importFileReader);
		for (String sql = reader.readLine(); sql != null; sql = reader
				.readLine()) {
			try {
				String trimmedSql = sql.trim();
				if (trimmedSql.length() == 0 || trimmedSql.startsWith("--")
						|| trimmedSql.startsWith("//")
						|| trimmedSql.startsWith("/*")) {
					continue;
				} else {
					if (trimmedSql.endsWith(";")) {
						trimmedSql = trimmedSql.substring(0, trimmedSql
								.length() - 1);
					}
					log.debug(trimmedSql);
					statement.execute(trimmedSql);
				}
			} catch (SQLException e) {
				throw new JDBCException("Error during import script execution",
						e);
			}
		}
	}

	private void create(boolean script, boolean export, Writer fileOutput,
			Statement statement) throws IOException {
		for (int j = 0; j < createSQL.length; j++) {
			try {
				execute(script, export, fileOutput, statement, createSQL[j]);
			} catch (SQLException e) {
				if (haltOnError) {
					throw new JDBCException("Error during DDL export", e);
				}
				exceptions.add(e);
				log.error("Unsuccessful: " + createSQL[j]);
				log.error(e.getMessage());
			}
		}
	}

	private void drop(boolean script, boolean export, Writer fileOutput,
			Statement statement) throws IOException {
		for (int i = 0; i < dropSQL.length; i++) {
			try {
				execute(script, export, fileOutput, statement, dropSQL[i]);
			} catch (SQLException e) {
				exceptions.add(e);
				log.debug("Unsuccessful: " + dropSQL[i]);
				log.debug(e.getMessage());
			}
		}
	}

	private void execute(boolean script, boolean export, Writer fileOutput,
			Statement statement, final String sql) throws IOException,
			SQLException {
		String formatted = format(sql);
		if (delimiter != null) {
			formatted += delimiter;
		}
		if (script) {
			System.out.println(formatted);
		}
		log.debug(formatted);
		if (outputFile != null) {
			fileOutput.write(formatted + "\n");
		}
		if (export) {
			statement.executeUpdate(sql);
			SQLWarning warnings = statement.getWarnings();
			if (warnings != null) {
				JDBCExceptionReporter.logAndClearWarnings(connectionHelper
						.getConnection());
			}
		}

	}

	public static void main(String[] args) {
		try {
			Configuration cfg = new Configuration();

			boolean script = true;
			boolean drop = false;
			boolean create = false;
			boolean halt = false;
			boolean export = true;
			String outFile = null;
			String importFile = "/import.sql";
			String propFile = null;
			boolean format = false;
			String delim = null;

			for (int i = 0; i < args.length; i++) {
				if (args[i].startsWith("--")) {
					if (args[i].equals("--quiet")) {
						script = false;
					} else if (args[i].equals("--drop")) {
						drop = true;
					} else if (args[i].equals("--create")) {
						create = true;
					} else if (args[i].equals("--haltonerror")) {
						halt = true;
					} else if (args[i].equals("--text")) {
						export = false;
					} else if (args[i].startsWith("--output=")) {
						outFile = args[i].substring(9);
					} else if (args[i].startsWith("--import=")) {
						importFile = args[i].substring(9);
					} else if (args[i].startsWith("--properties=")) {
						propFile = args[i].substring(13);
					} else if (args[i].equals("--format")) {
						format = true;
					} else if (args[i].startsWith("--delimiter=")) {
						delim = args[i].substring(12);
					} else if (args[i].startsWith("--config=")) {
						cfg.configure(args[i].substring(9));
					} else if (args[i].startsWith("--naming=")) {
						cfg.setNamingStrategy((NamingStrategy) ReflectHelper
								.classForName(args[i].substring(9))
								.newInstance());
					}
				} else {
					String filename = args[i];
					if (filename.endsWith(".jar")) {
						cfg.addJar(new File(filename));
					} else {
						cfg.addFile(filename);
					}
				}

			}

			if (propFile != null) {
				Properties props = new Properties();
				props.putAll(cfg.getProperties());
				props.load(new FileInputStream(propFile));
				cfg.setProperties(props);
			}

			SchemaExport se = new SchemaExport(cfg).setHaltOnError(halt)
					.setOutputFile(outFile).setImportFile(importFile)
					.setDelimiter(delim);
			if (format) {
				se.setFormat(true);
			}
			se.execute(script, export, drop, create);

		} catch (Exception e) {
			log.error("Error creating schema ", e);
			e.printStackTrace();
		}
	}

	/**
	 * Returns a List of all Exceptions which occured during the export.
	 * 
	 * @return A List containig the Exceptions occured during the export
	 */
	public List getExceptions() {
		return exceptions;
	}

	public MySchemaExport setHaltOnError(boolean haltOnError) {
		this.haltOnError = haltOnError;
		return this;
	}
}
